package ws

import (
	"context"
	"github.com/zeromicro/go-zero/core/executors"
	model "go-admin/base/mongo"
	"time"
)

const (
	flushInterval = time.Second
	maxBulkRows   = 1000
)

type MongoInserter struct {
	executor        *executors.PeriodicalExecutor
	mongoDBInserter *mongoDBInserter
}

func NewMongoInserter(conn model.MessageModel) (*MongoInserter, error) {
	mongoDBInserter := &mongoDBInserter{
		conn:     conn,
		maxSize:  maxBulkRows,
		messages: []model.Message{},
	}
	return &MongoInserter{
		executor:        executors.NewPeriodicalExecutor(flushInterval, mongoDBInserter),
		mongoDBInserter: mongoDBInserter,
	}, nil
}
func (m *MongoInserter) Flush() {
	m.executor.Flush()
}

func (m *MongoInserter) Insert(message model.Message) {
	m.executor.Add(message)
}

type mongoDBInserter struct {
	conn     model.MessageModel
	messages []model.Message
	maxSize  int64
}

func (m *mongoDBInserter) AddTask(task any) bool {
	m.messages = append(m.messages, task.(model.Message))
	return len(m.messages) >= maxBulkRows
}

func (m *mongoDBInserter) Execute(tasks any) {
	messages := tasks.([]model.Message)
	err := m.conn.InsertMany(context.Background(), messages)
	if err != nil {
		return
	}
}

func (m *mongoDBInserter) RemoveAll() any {
	values := m.messages
	m.messages = nil
	return values
}
